package main

import (
	"flag"
	"fmt"
	"gitee.com/quant1x/data/cache"
	"gitee.com/quant1x/data/category/trading"
	"gitee.com/quant1x/data/internal/runtimes"
	"gitee.com/quant1x/data/internal/tdx"
	"gitee.com/quant1x/data/security"
	"gitee.com/quant1x/data/stock"
	"gitee.com/quant1x/pandas"
	"gitee.com/quant1x/pandas/stat"
	"github.com/mymmsc/gox/cron"
	"github.com/mymmsc/gox/logger"
	"github.com/mymmsc/gox/progressbar"
	"github.com/mymmsc/gox/signal"
	"github.com/mymmsc/gox/util/treemap"
	"os"
	"sync"
	"time"
)

var (
	MinVersion string // 版本号
	today      bool   // 是否当日
)

// 更新tick数据
func main() {
	mainStart := time.Now()
	defer func() {
		elapsedTime := time.Since(mainStart) / time.Millisecond
		fmt.Printf("\n总耗时: %.3fs\n", float64(elapsedTime)/1000)
	}()
	//创建监听退出chan
	c := signal.Notify()
	var (
		cronConfig string // 定时脚本
		cronTrue   bool   // 是否启用应用内定时器
		startDate  string // 开始日期
		version    bool   // 显示版本号
	)
	flag.StringVar(&cronConfig, "cron_config", "0 0 17 * * ?", "pull code data cron")
	flag.BoolVar(&cronTrue, "cron_true", false, "use crontab in application")
	flag.StringVar(&startDate, "start", cache.TickStartDate, "tick start date")
	flag.BoolVar(&version, "version", false, "print version")
	flag.BoolVar(&today, "today", false, "want to update the current day's data")
	flag.Parse()

	if version {
		fmt.Println(MinVersion)
		os.Exit(0)
	}
	runtimes.GOMAXPROCS()
	if today && trading.CurrentlyTrading() {
		cache.UpdateTickStartDate(trading.LastTradeDate())
	} else {
		cache.UpdateTickStartDate(startDate)
	}
	if !cronTrue {
		handleCodeData()
	} else {
		crontab := cron.New(cron.WithSeconds()) //精确到秒
		// 添加定时任务,
		_, _ = crontab.AddFunc(cronConfig, handleCodeData)
		//启动定时器
		crontab.Start()
		select {
		case sig := <-c:
			{
				logger.Info("进程结束:", sig)
				os.Exit(1)
			}
		}
	}
}

func handleCodeData() {
	logger.Info("任务开始启动...")

	fullCodes := tdx.GetZxgList()
	fmt.Printf("自选股, 共计[%d]\n", len(fullCodes))
	if len(fullCodes) == 0 {
		fmt.Printf("没有指定自选股, 全量更新\n")
		fullCodes = stock.GetCodeList()
	}
	fmt.Printf("实时更新指数及个股, 共计[%d]\n", len(fullCodes))

	count := len(fullCodes)
	bar := progressbar.NewBar(1, "执行[更新历史tick数据]", count)
	var wg = sync.WaitGroup{}
	var mapKaiPanVol *treemap.Map
	mapKaiPanVol = treemap.NewWithStringComparator()
	for _, code := range fullCodes {
		basicInfo, err := security.GetBasicInfo(code)
		if err == security.ErrCodeNotExist {
			// 证券代码不存在
			bar.Add(1)
			continue
		}
		if err != nil {
			// 其它错误 输出错误信息
			logger.Errorf("%s => %v", code, err)
			bar.Add(1)
			continue
		}
		if basicInfo.Delisting {
			logger.Errorf("%s => 已退市", code)
			bar.Add(1)
			continue
		}
		wg.Add(1)
		pullData := func(pbar *progressbar.Bar, waitGroup *sync.WaitGroup, fullCode string) {
			defer waitGroup.Done()
			kpv := tdx.GetTickAll(fullCode)
			if today {
				mapKaiPanVol.Put(fullCode, kpv)
			}
			pbar.Add(1)
		}
		pullData(bar, &wg, code)
	}
	wg.Wait()
	codes := []string{}
	vols := []float64{}
	mapKaiPanVol.Each(func(key interface{}, value interface{}) {
		code, ok1 := key.(string)
		vol, ok2 := value.(float64)
		if ok1 && ok2 {
			codes = append(codes, code)
			vols = append(vols, vol)
		}
	})
	s1 := pandas.NewSeries(stat.SERIES_TYPE_STRING, "code", codes)
	s2 := pandas.NewSeries(stat.SERIES_TYPE_FLOAT64, "vol", vols)
	df := pandas.NewDataFrame(s1, s2)
	_ = df.WriteCSV(cache.TickKaiPanFilename())
	logger.Info("任务执行完毕.", time.Now())
}
